package com.ivideo.avcore.rtmplive;

import static android.media.MediaFormat.KEY_BIT_RATE;
import static android.media.MediaFormat.KEY_COLOR_FORMAT;
import static android.media.MediaFormat.KEY_FRAME_RATE;
import static android.media.MediaFormat.KEY_I_FRAME_INTERVAL;
import static android.media.MediaFormat.KEY_MAX_INPUT_SIZE;

import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaCodecList;
import android.media.MediaFormat;
import android.util.Log;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Date;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * @author：lxy
 * @decription: 媒体编码 封装视频和音频编码
 * @time：2025/7/2
 */
public class MediaEncoder {

    private static final String TAG = "MediaEncoder";
    private Config mConfig; //参数配置

    private MediaEncoder(Config config) {
        this.mConfig = config;
    }

    private volatile static MediaEncoder instance;

    public static MediaEncoder getInstance(Config config) {
        if (instance == null) {
            synchronized (MediaEncoder.class) {
                if (instance == null) {
                    instance = new MediaEncoder(config);
                }
            }
        }
        return instance;
    }

    private long presentationTimeUs;
    private boolean videoEncoderLoop, audioEncoderLoop;
    private Thread videoEncoderThread, audioEncoderThread;
    private MediaCodec vEncoder; //视频编码器
    private MediaCodec aEncoder; //音频编码器

    private LinkedBlockingQueue<byte[]> videoQueue; //视频数据队列：摄像头采集后的数据放到此队列中
    private LinkedBlockingQueue<byte[]> audioQueue;//音频数据队列：AudidoRecoder采集后的数据放到此队列中
    private MediaCodec.BufferInfo vBufferInfo = new MediaCodec.BufferInfo(),
            aBufferInfo = new MediaCodec.BufferInfo(); //编码缓存

    /**
     * 添加视频数据
     *
     * @param data
     */
    public void putVideoData(byte[] data) {
        try {
            videoQueue.put(data);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 添加音频数据
     *
     * @param data
     */
    public void putAudioData(byte[] data) {
        try {
            audioQueue.put(data);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }


    /**
     * 开始编码
     */
    public void start() {
        startAudioEncode();
        startVideoEncode();
    }

    /**
     * 停止编码
     */
    public void stop() {
        stopAudioEncode();
        stopVideoEncode();

    }

    /**
     * 释放资源
     */
    public void release() {
        releaseAudioEncoder();
        releaseVideoEncoder();
    }

    private void releaseVideoEncoder() {
        if (vEncoder != null) {
            vEncoder.release();
            vEncoder = null;
        }
    }

    private void releaseAudioEncoder() {
        if (aEncoder != null) {
            aEncoder.release();
            aEncoder = null;
        }

    }

    /**
     * 停止音频编码
     */
    private void stopAudioEncode() {
        audioEncoderLoop = false;
        if (audioEncoderThread != null) {
            audioEncoderThread.interrupt();
        }
        if (aEncoder != null) {
            aEncoder.stop();
        }
    }

    /**
     * 停止视频编码
     */
    private void stopVideoEncode() {
        videoEncoderLoop = false;
        if (videoEncoderThread != null) {
            videoEncoderThread.interrupt();
        }
        if (vEncoder != null) {
            vEncoder.stop();
        }
    }

    private void startVideoEncode() {
        if (vEncoder == null) {
            throw new RuntimeException("请初始化视频编码器");
        }

        if (videoEncoderLoop) {
            throw new RuntimeException("必须先停止");
        }
        videoEncoderThread = new Thread(new Runnable() {
            @Override
            public void run() {
                presentationTimeUs = System.currentTimeMillis() * 1000;
                vEncoder.start();
                while (videoEncoderLoop && !Thread.interrupted()) {
                    try {
                        byte[] data = videoQueue.take(); //取出待编码数据
                        encodeVideoData(data); //视频编码
                    } catch (InterruptedException e) {
                        throw new RuntimeException(e);
                    }

                }


            }
        });
        videoEncoderLoop = true;
        videoEncoderThread.start();
    }


    public void startAudioEncode() {
        if (aEncoder == null) {
            throw new RuntimeException("请初始化音频编码器");
        }

        if (audioEncoderLoop) {
            throw new RuntimeException("必须先停止");
        }
        audioEncoderThread = new Thread(new Runnable() {
            @Override
            public void run() {
                presentationTimeUs = System.currentTimeMillis() * 1000;
                aEncoder.start();
                while (audioEncoderLoop && !Thread.interrupted()) {
                    try {
                        byte[] data = audioQueue.take(); //取出待编码数据
                        encodeAudioData(data); //音频编码
                    } catch (InterruptedException e) {
                        throw new RuntimeException(e);
                    }

                }


            }
        });
        audioEncoderLoop = true;
        audioEncoderThread.start();

    }

    /**
     * 视频编码
     *
     * @param data
     */
    private void encodeVideoData(byte[] data) {
        ByteBuffer[] inputBuffers = vEncoder.getInputBuffers();
        ByteBuffer[] outputBuffers = vEncoder.getOutputBuffers();
        //喂数据
        int inputBufferId = vEncoder.dequeueInputBuffer(-1);
        if (inputBufferId >= 0) {
            ByteBuffer bb = inputBuffers[inputBufferId];
            bb.clear();
            bb.put(data, 0, data.length);
            long pts = new Date().getTime() * 1000 - presentationTimeUs;
            vEncoder.queueInputBuffer(inputBufferId, 0, data.length, pts, 0);
        }
        //取数据：编码后数据---》回调外界
        int outputBufferId = vEncoder.dequeueOutputBuffer(vBufferInfo, 0);
        if (outputBufferId >= 0) {
            ByteBuffer bb = outputBuffers[outputBufferId];
            if (mCallback != null) {
                mCallback.outputVideoData(bb, vBufferInfo);
            }
            vEncoder.releaseOutputBuffer(outputBufferId, false);
        }
    }

    /**
     * 音频编码
     *
     * @param data
     */
    public void encodeAudioData(byte[] data) {
        ByteBuffer[] inputBuffers = aEncoder.getInputBuffers();
        ByteBuffer[] outputBuffers = aEncoder.getOutputBuffers();
        //喂数据
        int inputBufferId = aEncoder.dequeueInputBuffer(-1);
        if (inputBufferId >= 0) {
            ByteBuffer bb = inputBuffers[inputBufferId];
            bb.clear();
            bb.put(data, 0, data.length);
            long pts = new Date().getTime() * 1000 - presentationTimeUs;
            aEncoder.queueInputBuffer(inputBufferId, 0, data.length, pts, 0);
        }
        //取数据：编码后数据---》回调外界
        int outputBufferId = aEncoder.dequeueOutputBuffer(aBufferInfo, 0);
        if (outputBufferId >= 0) {
            ByteBuffer bb = outputBuffers[outputBufferId];
            if (mCallback != null) {
                mCallback.outputAudioData(bb, aBufferInfo);
            }
            aEncoder.releaseOutputBuffer(outputBufferId, false);
        }

    }


    /**
     * 初始化音频编码器
     *
     * @param sampleRate
     * @param channelCount
     * @return
     * @throws IOException
     */
    public void initAudioEncoder(int sampleRate, int channelCount) throws IOException {
        MediaCodec encoder = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_AUDIO_AAC);
        MediaFormat format = MediaFormat.createAudioFormat(MediaFormat.MIMETYPE_AUDIO_AAC, sampleRate, channelCount);
        format.setInteger(KEY_MAX_INPUT_SIZE, 0);
        format.setInteger(KEY_BIT_RATE, sampleRate * channelCount);
        encoder.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
        aEncoder = encoder;
        audioQueue = new LinkedBlockingQueue<>();
    }


    /**
     * 初始化视频编码器
     *
     * @param width
     * @param height
     * @return
     * @throws IOException
     */
    public int initVideoEncoder(int width, int height) throws IOException {
        //获得当前设备支持的编码器
        MediaCodecInfo mediaCodecInfo = getMediaCodecInfoByType(MediaFormat.MIMETYPE_VIDEO_AVC);
        //获得编码器支持的颜色格式
        int colorFormat = getColorFormat(mediaCodecInfo);
        //创建编码器
        MediaCodec mediaCodec = MediaCodec.createByCodecName(mediaCodecInfo.getName());
        //配置编码器的参数
        MediaFormat format = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, width, height);
        format.setInteger(KEY_MAX_INPUT_SIZE, 0);
        format.setInteger(KEY_BIT_RATE, mConfig.bitrate); //码率
        format.setInteger(KEY_COLOR_FORMAT, colorFormat);//颜色格式
        format.setInteger(KEY_FRAME_RATE, mConfig.fps); //帧率
        format.setInteger(KEY_I_FRAME_INTERVAL, mConfig.iFrameInternal); //关键帧间隔
        mediaCodec.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);//CONFIGURE_FLAG_ENCODE编码
        vEncoder = mediaCodec;
        videoQueue = new LinkedBlockingQueue<>();
        return colorFormat;
    }

    /**
     * 根据编码器信息获得格式
     *
     * @param mediaCodecInfo
     * @return
     */
    private static int getColorFormat(MediaCodecInfo mediaCodecInfo) {
        int matchedForamt = 0;
        MediaCodecInfo.CodecCapabilities codecCapabilities =
                mediaCodecInfo.getCapabilitiesForType(MediaFormat.MIMETYPE_VIDEO_AVC);
        for (int i = 0; i < codecCapabilities.colorFormats.length; i++) {
            int format = codecCapabilities.colorFormats[i];
            if (format >= MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar &&
                    format <= MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420PackedSemiPlanar
            ) {
                if (format >= matchedForamt) {
                    matchedForamt = format;
                    break;
                }
            }
        }
        switch (matchedForamt) {
            case MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar:
                Log.i(TAG, "selected yuv420p");
                break;
            case MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420PackedPlanar:
                Log.i(TAG, "selected yuv420pp");
                break;
            case MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar:
                Log.i(TAG, "selected yuv420sp");
                break;
            case MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420PackedSemiPlanar:
                Log.i(TAG, "selected yuv420psp");
                break;

        }
        return matchedForamt;
    }

    private static MediaCodecInfo getMediaCodecInfoByType(String mimeType) {
        for (int i = 0; i < MediaCodecList.getCodecCount(); i++) {
            MediaCodecInfo codecInfo = MediaCodecList.getCodecInfoAt(i);
            if (!codecInfo.isEncoder()) {
                continue;
            }
            String[] types = codecInfo.getSupportedTypes();
            for (int j = 0; j < types.length; j++) {
                if (types[j].equalsIgnoreCase(mimeType)) {
                    return codecInfo;
                }
            }
        }
        return null;
    }

    private Callback mCallback;

    /**
     * 设置回调
     *
     * @param callback 回调
     */
    public void setCallback(Callback callback) {
        this.mCallback = callback;
    }


    public interface Callback {
        void outputVideoData(ByteBuffer bb, MediaCodec.BufferInfo info);

        void outputAudioData(ByteBuffer bb, MediaCodec.BufferInfo info);
    }
}
